import os
import torch
from PIL import Image
from torch.utils.data import Dataset
from utils import label2class


class Flowers(Dataset):
    # 用于读取flower数据集
    def __init__(self, dataset_path: str, transforms=None):
        '''
        存储所有数据 data路径和label
        :param dataset_path:
        '''
        super(Dataset, self).__init__()
        flowers = os.listdir(dataset_path)
        flowers = sorted(flowers) # 必须排序,否在每一次顺序不一样训练测试类别就会乱
        self.flower_paths = []
        self.class2label = {}  # 类别str 转 label
        label = 0
        for _, flower in enumerate(flowers):
            flowers_path = os.path.join(dataset_path, flower)
            if os.path.isdir(flowers_path):
                self.class2label[flower] = label
                label +=1
                sub_flowers = os.listdir(flowers_path)
                for sub_flower in sub_flowers:
                    self.flower_paths.append(os.path.join(flowers_path, sub_flower))
        self.label2class = label2class(self.class2label)  # label 转 类别str
        self.transforms = transforms

    def __len__(self):
        return len(self.flower_paths)

    def __getitem__(self, item):
        # 读取数据和label
        img = Image.open(self.flower_paths[item])
        label = self.class2label[self.flower_paths[item].split('/')[-2]]
        if self.transforms is not None:
            img = self.transforms(img)  # 数据增强
        return img, label

    @staticmethod
    def collate_fn(batch):
        # 官方实现的default_collate可以参考
        # https://github.com/pytorch/pytorch/blob/67b7e751e6b5931a9f45274653f4f653a4e6cdf6/torch/utils/data/_utils/collate.py
        images, labels = tuple(zip(*batch))

        images = torch.stack(images, dim=0)
        labels = torch.as_tensor(labels)
        return images, labels


# 通过可视化测试数据读取是否正常
if __name__ == '__main__':
    import matplotlib.pyplot as plt

    flowers = Flowers(dataset_path='../datasets/flower_photos-mini')
    img, label = flowers.__getitem__(10)
    plt.imshow(img)
    plt.title(flowers.label2class[label])
    plt.show()
    pass
